from lib.conn import Conn
import pandas as pd
from lib.func import *
import talib as ta


def daily_idc(ts_code, start, end=None):
    conn = Conn()
    if not end:
        end = today()
    start_year = table_suffix(start)
    end_year = table_suffix(end)
    field_str = 'ts_code,trade_date,high,low,close,`change`,vol,amount'
    if start_year == end_year:
        sql = 'select {} from td_market_daily_{} where ts_code=%s and trade_date between %s and %s order by trade_date'.format(field_str, start_year)
        arr = conn.fetch_all(sql, [ts_code, start, end])
    else:
        sql = '(select {} from td_market_daily_{} where ts_code=%s and trade_date between %s and %s) ' \
              'union all ' \
              '(select {} from td_market_daily_{} where ts_code=%s and trade_date between %s and %s) ' \
              'ORDER BY trade_date'.\
            format(field_str, start_year, field_str, end_year)
        arr = conn.fetch_all(sql, [ts_code, start, start_year + '1231', ts_code, end_year + '0101', end])
    df = pd.DataFrame(arr)
    if len(arr) == 0:
        return df
    return calc_magic(df)


def table_suffix(date):
    return date[:4]


def week_idc(ts_code, end=None):
    conn = Conn()
    if not end:
        end = today()
    sql = 'select ts_code,trade_date,high,low,close,`change`,vol,amount from td_market_weekly ' \
          'where ts_code=%s and trade_date<=%s order by trade_date'
    arr = conn.fetch_all(sql, [ts_code, end])
    df = pd.DataFrame(arr)
    if len(arr) == 0:
        return df
    try:
        df.set_index("trade_date", inplace=True)
    except KeyError:
        print(arr)
    df.sort_index()
    df = merge_weekly(df)
    return df


def calc_magic(df):
    df['ma_long'] = ta.MA(df.close, timeperiod=100, matype=0)
    df['ma_middle'] = ta.MA(df.close, timeperiod=75, matype=0)
    df['lm_diff'] = df['ma_middle'] - df['ma_long']
    df['lm_diff_change'] = df['lm_diff'].diff(periods=1)
    df['ma_short'] = ta.MA(df.close, timeperiod=25, matype=0)
    df['ma_long_change'] = df['ma_long'].diff(periods=1)
    df['rsi'] = ta.RSI(df.close, timeperiod=9)
    df['strong'] = df['vol'] * df['change']
    df['strong_short'] = df['strong'].ewm(span=2, adjust=False).mean()
    df['strong_long'] = df['strong'].ewm(span=13, adjust=False).mean()
    return df


def calc_macd(df):
    df['ema_short'] = df['close'].ewm(span=12, adjust=False).mean()
    df['ema_long'] = df['close'].ewm(span=26, adjust=False).mean()
    # 计算DIFF、DEA、MACD
    df['diff'] = df['ema_short'] - df['ema_long']
    df['dea'] = df['diff'].ewm(span=9, adjust=False).mean()
    df['macd'] = 2 * (df['diff'] - df['dea'])
    return df


def merge_weekly(df):
    # 计算EMA(12)和EMA(26)
    # 计算复权open、high、low
    df['ema_short'] = df['close'].ewm(span=12, adjust=False).mean()
    df['ema_long'] = df['close'].ewm(span=26, adjust=False).mean()
    # 计算DIFF、DEA、MACD
    df['diff'] = df['ema_short'] - df['ema_long']
    df['dea'] = df['diff'].ewm(span=9, adjust=False).mean()
    df['macd'] = 2 * (df['diff'] - df['dea'])
    df['macd_change'] = df['macd'].diff(periods=1)
    df['ema_change'] = df['ema_short'].diff(periods=1)
    return df


def merge_daily(df):
    # 计算强力指标
    df['fi'] = df['vol'] * df['change']
    df['fi_2'] = df['fi'].ewm(span=2, adjust=False).mean()
    return df

